# yolo11s/yolo_wrapper.py
import os
import numpy as np

try:
    from ultralytics import YOLO
except Exception as e:
    YOLO = None
    _YOLO_IMPORT_ERROR = e
else:
    _YOLO_IMPORT_ERROR = None


# COCO id -> name
COCO_NAMES = {
    0: 'person', 1: 'bicycle', 2: 'car', 3: 'motorcycle', 4: 'airplane', 5: 'bus', 6: 'train', 7: 'truck',
    8: 'boat', 9: 'traffic light', 10: 'fire hydrant', 11: 'stop sign', 12: 'parking meter', 13: 'bench',
    14: 'bird', 15: 'cat', 16: 'dog', 17: 'horse', 18: 'sheep', 19: 'cow', 20: 'elephant', 21: 'bear',
    22: 'zebra', 23: 'giraffe', 24: 'backpack', 25: 'umbrella', 26: 'handbag', 27: 'tie', 28: 'suitcase',
    29: 'frisbee', 30: 'skis', 31: 'snowboard', 32: 'sports ball', 33: 'kite', 34: 'baseball bat',
    35: 'baseball glove', 36: 'skateboard', 37: 'surfboard', 38: 'tennis racket', 39: 'bottle',
    40: 'wine glass', 41: 'cup', 42: 'fork', 43: 'knife', 44: 'spoon', 45: 'bowl', 46: 'banana',
    47: 'apple', 48: 'sandwich', 49: 'orange', 50: 'broccoli', 51: 'carrot', 52: 'hot dog', 53: 'pizza',
    54: 'donut', 55: 'cake', 56: 'chair', 57: 'couch', 58: 'potted plant', 59: 'bed', 60: 'dining table',
    61: 'toilet', 62: 'tv', 63: 'laptop', 64: 'mouse', 65: 'remote', 66: 'keyboard', 67: 'cell phone',
    68: 'microwave', 69: 'oven', 70: 'toaster', 71: 'sink', 72: 'refrigerator', 73: 'book', 74: 'clock',
    75: 'vase', 76: 'scissors', 77: 'teddy bear', 78: 'hair drier', 79: 'toothbrush'
}
NAME_TO_ID = {v: k for k, v in COCO_NAMES.items()}


class YoloSegmentor:
    """
    用 YOLO 分割，然后把深度图上不是指定类的像素清零。
    支持：
    - target 用名字 ('bottle') 或 id (39)
    - YOLO 输出的 mask 是 2D / 3D / 4D 都能吃
    - 自动 resize 到和深度一样大
    """

    def __init__(self, model_path: str, device=None,
                 output_dir: str = '/home/spaceai/anygrasp_sdk/yolo11s/output'):
        if YOLO is None:
            raise RuntimeError(f'ultralytics 未安装: {_YOLO_IMPORT_ERROR}')
        self.model = YOLO(model_path)
        self.device = device
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)

        self.last_coverage = 0.0
        self.last_overlay_path = None
        self.last_mask_path = None

    # ============== 工具函数 ==============

    def _to_uint8_rgb(self, img: np.ndarray) -> np.ndarray:
        if img.dtype.kind == 'f':
            return (img * 255.0).clip(0, 255).astype(np.uint8)
        return img.astype(np.uint8)

    def _resolve_target_to_id(self, target, model_names: dict):
        # int 直接用
        if isinstance(target, int):
            return target
        if not isinstance(target, str):
            return None

        # 先查模型自己的 names
        if model_names:
            for k, v in model_names.items():
                if v == target:
                    return int(k)

        # 再查我们这张 COCO 表
        return NAME_TO_ID.get(target, None)

    def _save_overlay(self, color_u8, mask_bool,
                      save_overlay_path=None, save_mask_path=None):
        self.last_overlay_path = None
        self.last_mask_path = None

        if color_u8 is None:
            return

        h, w = color_u8.shape[:2]

        if save_overlay_path is None:
            save_overlay_path = os.path.join(self.output_dir, 'last_yolo_overlay.png')
        if save_mask_path is None:
            save_mask_path = os.path.join(self.output_dir, 'last_yolo_mask.png')

        try:
            import cv2
            overlay = color_u8.copy()
            if mask_bool is not None:
                # 尺寸不对的话再拉一层
                if mask_bool.shape != (h, w):
                    m = (mask_bool.astype('uint8') * 255)
                    m = cv2.resize(m, (w, h), interpolation=cv2.INTER_NEAREST)
                    mask_vis = m > 0
                else:
                    mask_vis = mask_bool

                alpha = 0.5
                red = np.zeros_like(overlay)
                red[:, :, 2] = 255
                overlay[mask_vis] = (
                    overlay[mask_vis].astype(np.float32) * (1 - alpha)
                    + red[mask_vis].astype(np.float32) * alpha
                ).astype(np.uint8)

                cv2.imwrite(save_mask_path, (mask_bool.astype('uint8') * 255))
                self.last_mask_path = save_mask_path

            cv2.imwrite(save_overlay_path, cv2.cvtColor(overlay, cv2.COLOR_RGB2BGR))
            self.last_overlay_path = save_overlay_path
            print(f'[YoloSegmentor] saved overlay to: {save_overlay_path}')
        except Exception as e:
            print('[YoloSegmentor] failed to save overlay/mask:', e)

    # 关键：把各种形状的 mask 全部变成 (N, H, W)
    def _normalize_masks(self, masks_np: np.ndarray):
        """
        把 YOLO 返回的 mask 统一成 (N, H, W)
        常见几种情况：
        - (H, W)                  -> (1, H, W)
        - (N, H, W)               -> 原样
        - (1, N, H, W)            -> squeeze(0)
        - (N, 1, H, W)            -> squeeze(1)
        """
        if masks_np.ndim == 2:
            return masks_np[None, :, :]              # (H, W) -> (1, H, W)

        if masks_np.ndim == 3:
            return masks_np                           # (N, H, W)

        if masks_np.ndim == 4:
            # 你遇到的就是这个： (1, 3, 384, 640)
            if masks_np.shape[0] == 1:
                masks_np = masks_np[0]                # -> (3, H, W)
                if masks_np.ndim == 3:
                    return masks_np
            if masks_np.shape[1] == 1:
                masks_np = masks_np[:, 0, :, :]       # -> (N, H, W)
                return masks_np
            # 再不行就直接把前两维合并
            n1, n2, h, w = masks_np.shape
            return masks_np.reshape(n1 * n2, h, w)

        # 非预期，直接返回空
        return np.zeros((0, 1, 1), dtype=np.uint8)

    # ============== 主函数（你主程序里调的） ==============

    def mask_depth(self,
                   color_img: np.ndarray,
                   depth_img: np.ndarray,
                   target='bottle',
                   conf: float = 0.25,
                   save_overlay_path: str = None):

        self.last_coverage = 0.0

        # 深度图尺寸 = 我们最终要对齐到的尺寸
        H, W = depth_img.shape[:2]

        # 1. 颜色转 uint8
        color_u8 = self._to_uint8_rgb(color_img)

        # 2. YOLO 推理：优先试试用和深度一样的尺寸；不行再 fallback
        #    这样你就能看到那个警告：720 -> 736
        # try:
        #     infer_kwargs = dict(imgsz=[H, W], conf=conf, verbose=False)
        #     if self.device is not None:
        #         infer_kwargs['device'] = self.device
        #     results = self.model(color_u8, imgsz=[H, W])
        # except TypeError:
        #     # 有些版本不认 imgsz=[H,W]，再试一次只给高
        #     try:
        #         results = self.model(color_u8, imgsz=H, conf=conf, verbose=False)
        #     except Exception:
        #         results = self.model(color_u8)
        # except Exception:
        #     # 最兜底
        results = self.model(color_u8)

        if len(results) == 0:
            self._save_overlay(color_u8, None, save_overlay_path)
            return depth_img

        res = results[0]

        # 3. 没有 masks 就原样返回
        if getattr(res, 'masks', None) is None:
            self._save_overlay(color_u8, None, save_overlay_path)
            return depth_img

        # 4. 取掩码数据
        try:
            masks = res.masks.data
        except Exception:
            try:
                masks = res.masks.masks
            except Exception:
                self._save_overlay(color_u8, None, save_overlay_path)
                return depth_img

        # -> numpy
        if hasattr(masks, 'cpu'):
            masks_np = masks.cpu().numpy()
        else:
            masks_np = np.array(masks)

        # 统一成 (N, Hm, Wm)
        masks_np = self._normalize_masks(masks_np)

        # 5. 取类别 id
        cls_ids = None
        if getattr(res, 'boxes', None) is not None and getattr(res.boxes, 'cls', None) is not None:
            try:
                cls_ids = res.boxes.cls.cpu().numpy().astype(int)
            except Exception:
                cls_ids = np.array(res.boxes.cls).astype(int)

        # 6. 把 target 转成真正的 id
        class_id = self._resolve_target_to_id(target, getattr(res, 'names', None))
        if class_id is None:
            self._save_overlay(color_u8, None, save_overlay_path)
            return depth_img

        # 7. 合成一张跟深度一样大的 keep_mask
        keep_mask = np.zeros((H, W), dtype=np.uint8)

        for i in range(masks_np.shape[0]):
            use = True
            if cls_ids is not None and i < len(cls_ids):
                use = (int(cls_ids[i]) == int(class_id))
            else:
                # 没拿到 cls，就尝试名字兜底
                try:
                    model_names = res.names
                    if isinstance(model_names, dict):
                        if getattr(res.boxes, 'cls', None) is not None and i < len(res.boxes.cls):
                            cid_i = int(res.boxes.cls[i])
                            use = (model_names.get(cid_i, '') == COCO_NAMES.get(class_id, ''))
                except Exception:
                    pass

            if not use:
                continue

            m = masks_np[i]          # (Hm, Wm)
            # YOLO 出来一般是 float 0~1
            m = (m * 255).astype(np.uint8)

            # 核心：这里统一 resize 到你的深度大小 (W, H)
            if m.shape != (H, W):
                try:
                    import cv2
                    m = cv2.resize(m, (W, H), interpolation=cv2.INTER_NEAREST)
                except Exception:
                    from PIL import Image
                    m = np.array(Image.fromarray(m).resize((W, H)))
            keep_mask = np.maximum(keep_mask, m)

        # 8. 没分到指定类
        if keep_mask.max() == 0:
            self._save_overlay(color_u8, None, save_overlay_path)
            return depth_img

        keep_bool = keep_mask > 0

        # 记录覆盖率
        self.last_coverage = float(keep_bool.sum()) / float(H * W)
        print(f'[YoloSegmentor] mask coverage={self.last_coverage * 100:.2f}%')

        # 保存可视化
        self._save_overlay(color_u8, keep_bool, save_overlay_path)

        # 9. 应用到深度图
        out_depth = np.array(depth_img, copy=True)
        if out_depth.dtype.kind == 'f':
            out_depth[~keep_bool] = 0.0
        else:
            out_depth[~keep_bool] = 0
        return out_depth
